import backtrader as bt
import requests
import datetime
import time
import tushare as ts
import pandas as pd


class SimpleLivingData(bt.feed.DataBase):
    params = (
        ('sleeptime', 1),
        ('symbol', 'ixic'),
        ('startday', ''),
        ('datasource', 'living')
    )

    def __init__(self, symbol, **kwargs):
        '''symbol: ixic hkHSI'''
        super().__init__(**kwargs)
        self.move_data = None
        self.tick = None
        self.tick_shift = 0
        self.symbol = symbol
        # self.sleeptime = symbol
        # used to stop feed when no more new data from web
        self.last_dt = datetime.datetime.now()

    def start(self):
        super().start()
        # add filter adjust time
        self.resample(timeframe=self.p.timeframe, compression=self.p.compression)

    def stop(self):
        pass

    def islive(self):
        ''' tell cerebro this is a living feed'''
        return True

    def _load(self):
        # request data
        msg = self._req_data()
        if msg is None:
            # No more data
            return False

        # fill the lines
        self.lines.datetime[0] = self.date2num(msg['datetime'])
        self.lines.open[0] = msg['open']
        self.lines.high[0] = msg['high']
        self.lines.low[0] = msg['low']
        self.lines.close[0] = msg['close']
        self.lines.volume[0] = msg['volume']
        self.lines.openinterest[0] = -1

        # Say success
        return True

    def _req_data(self):
        '''internal function used to request data from the web
            base on the self.symbol it will fetch different data source'''

        # sleep some time
        time.sleep(self.p.sleeptime)

        msg = dict()
        if self.symbol == 'ixic':
            # ixic
            ret = requests.get('https://hq.sinajs.cn/list=gb_ixic')
            lstr = ret.text.split(',')

            dt = datetime.datetime.strptime(lstr[3], '%Y-%m-%d %H:%M:%S')
            if self.last_dt == dt:
                return None
            else:
                self.last_dt = dt
            msg['datetime'] = dt
            msg['open'] = float(lstr[5])
            msg['high'] = float(lstr[6])
            msg['low'] = float(lstr[7])
            msg['close'] = float(lstr[1])
            msg['volume'] = int(lstr[10])

        elif self.symbol == 'hkHSI':
            # hkHSI
            ret = requests.get('https://hq.sinajs.cn/list=rt_hkHSI')
            lstr = ret.text.split(',')

            dt = datetime.datetime.strptime(lstr[17] + ' ' + lstr[18], '%Y/%m/%d %H:%M:%S')
            if self.last_dt == dt:
                return None
            else:
                self.last_dt = dt
            msg['datetime'] = dt
            msg['open'] = float(lstr[2])
            msg['high'] = float(lstr[4])
            msg['low'] = float(lstr[5])
            msg['close'] = float(lstr[6])
            msg['volume'] = int(float(lstr[11]))
            print(msg)
        elif self.symbol == 'szzs':
            # hkHSI
            ret = requests.get('http://hq.sinajs.cn/list=sh000001')  # 5秒更新
            lstr = ret.text.split(',')
            # print(lstr)
            dt = datetime.datetime.strptime(lstr[30] + ' ' + lstr[31], '%Y-%m-%d %H:%M:%S')
            if self.last_dt == dt:
                return None
            else:
                self.last_dt = dt
            msg['datetime'] = dt
            msg['open'] = float(lstr[1])
            msg['high'] = float(lstr[4])
            msg['low'] = float(lstr[5])
            msg['close'] = float(lstr[3])
            msg['volume'] = int(float(lstr[9]))
        elif self.symbol == 'nyhj':
            # hkHSI
            ret = requests.get('http://hq.sinajs.cn/list=hf_GC')  # 6秒更新
            lstr = ret.text.split(',')
            dt = datetime.datetime.strptime(lstr[12] + ' ' + lstr[6], '%Y-%m-%d %H:%M:%S')
            # if self.last_dt == dt:
            #     return None
            # else:
            #     self.last_dt = dt
            msg['datetime'] = dt
            msg['open'] = float(lstr[2])
            msg['high'] = float(lstr[4])
            msg['low'] = float(lstr[5])
            msg['close'] = float(lstr[3])
            msg['volume'] = int(float(lstr[9]))
        elif self.p.datasource == 'day_sim':  # 测试代码，根据回测日期一天一天提供数据
            if (self.move_data == None):
                self.move_data = datetime.date(2021, 1, 1)  # datetime.datetime.now() - datetime.timedelta(days = 5) #
            # self.move_data = datetime.datetime.now() - datetime.timedelta(days=5)  # datetime(2020, 1, 1)
            self.move_data_start_string = (self.move_data-datetime.timedelta(days=10)).strftime('%Y-%m-%d')
            self.move_data_end_string = self.move_data.strftime('%Y-%m-%d')
            self.move_data = self.move_data + datetime.timedelta(days=1)

            print("tushare sim data")
            df = ts.get_k_data(self.symbol, autype='qfq', start=self.move_data_start_string,
                               end=self.move_data_end_string).sort_index(ascending=False)
            df.index = pd.to_datetime(df.date)
            dt = datetime.datetime.strptime(str(df.date[0]) + ' ' + str("00:10:10"), '%Y-%m-%d %H:%M:%S')
            msg['datetime'] = dt
            msg['open'] = float(df.open[0])
            msg['high'] = float(df.high[0])
            msg['low'] = float(df.low[0])
            msg['close'] = float(df.close[0])
            msg['volume'] = int(df.volume[0])
        elif self.p.datasource == 'tick_sim':  # 测试代码，根据回测tushare tick提供数据
            if self.tick is None:
                # self.tick = ts.get_today_ticks(self.symbol)
                self.tick = ts.get_tick_data(self.symbol,date=str((datetime.datetime.today()).date()),src='tt')
                self.tick_index = 0
            print("tushare tick sim data")
            df = self.tick.iloc[self.tick_index]
            self.tick_index = self.tick_index+1
            # df.index = pd.to_datetime(df.time)
            msg['datetime'] = datetime.datetime.strptime("9999-02-22" + ' ' + df.time, '%Y-%m-%d %H:%M:%S')
            msg['open'] = float(df.price)
            msg['high'] = float(df.price)
            msg['low'] = float(df.price)
            msg['close'] = float(df.price)
            msg['volume'] = int(df.amount)
        else:
            df = ts.get_realtime_quotes(self.symbol)
            df.index = pd.to_datetime(df.date)
            # print("tushare real time data")
            dt = datetime.datetime.strptime(str(df.date[0]) + ' ' + str(df.time[0]), '%Y-%m-%d %H:%M:%S')
            msg['datetime'] = dt
            msg['open'] = float(df.open[0])
            msg['high'] = float(df.high[0])
            msg['low'] = float(df.low[0])
            msg['close'] = float(df.price[0])
            msg['volume'] = int(df.volume[0])

        # print(msg)
        return msg
